# image vector db

import chromadb
import gradio as gr
import requests
from chromadb.utils import embedding_functions


class OpenAIEmbeddingFunction(embedding_functions.EmbeddingFunction):
    def __init__(self, model_name="bge-m3",
                 embedding_url="http://127.0.0.1:11434/v1/embeddings"):
        self.model_name = model_name
        self.embedding_url = embedding_url
        self.max_text_length = 768

    def __call__(self, text):
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer Empty"
        }
        if len(text) >= self.max_text_length:
            text = text[:self.max_text_length - 1]
        response = requests.post(self.embedding_url,
                                 headers=headers,
                                 json={"model": self.model_name, "input": [text]})
        data = response.json()["data"][0]["embedding"]
        return data


def semantic_search(query, embedding_function, db_path="vector_db", top_k=5):
    """基于语义的搜索"""
    client = chromadb.PersistentClient(path=db_path)
    collection = client.get_collection(name="image_descriptions")
    query_embedding = embedding_function(query)

    results = collection.query(
        query_embeddings=query_embedding,
        n_results=top_k
    )
    return results


def search_images_by_query(query):
    embedding_function = OpenAIEmbeddingFunction(model_name="bge-m3:latest",
                                                 embedding_url="http://127.0.0.1:11434/v1/embeddings")
    db_path = "../rag/vector_db"

    # 执行语义搜索
    results = semantic_search(query, embedding_function, db_path=db_path)

    # 解析结果，提取图像路径和描述
    images_and_descriptions = []
    for idx, doc in enumerate(results['documents'][0]):
        image_path = results['metadatas'][0][idx]['image_path']
        description = doc
        images_and_descriptions.append((image_path, description))

    return images_and_descriptions


# 创建Gradio界面
with gr.Blocks() as demo:
    gr.Markdown("# 图像知识库检索系统")
    with gr.Row():
        query_input = gr.Textbox(label="输入查询", placeholder="请输入关于你想要找的图片的描述...")
    with gr.Row():
        search_button = gr.Button("搜索")
    with gr.Row():
        output_gallery = gr.Gallery(label="搜索结果")

    def on_search_click(query):
        images_and_descriptions = search_images_by_query(query)
        gallery_items = [(img_desc[0], img_desc[1]) for img_desc in images_and_descriptions]
        return gallery_items

    search_button.click(on_search_click, inputs=query_input, outputs=output_gallery)

demo.launch(allowed_paths=[r"C:\WorkSpace\jingxiang.ai\aipcagent\examples\rag_data\images"])
